/******************************************************
 * @file get_weights.c
 * @author Destiny 
 * @brief 模型参数读取相关API定义
 * 
 * @version 0.1
 * @date 2024-07-18
 *****************************************************/
#include <stdio.h>
#include <stdlib.h>
#include <hdf5.h>
#include <string.h>
#include "cnn_operation.h"
#include "model_manager.h"



static hid_t g_File_id;   //模型参数文件FD

/**
 * @brief 为ptKernelParams分配内存（malloc）
 * 
 * @param ptKernelParams 模型参数
 */
static void MallocForParams(PT_KernelParams ptKernelParams)
{
    int i, j, k;
    ptKernelParams->ppppfKernelParams = (float ****)malloc(sizeof(float ***) * ptKernelParams->iRow);
    while (ptKernelParams->ppppfKernelParams == NULL)
    {
        DBG_PRINTF(DLOG_LVL_ERROR,DLOG_TAG,"ppppfKernelParams[] malloc  faild!\r\n");
        ptKernelParams->ppppfKernelParams = (float ****)malloc(sizeof(float ***) * ptKernelParams->iRow);
    }
    for (k = 0; k < ptKernelParams->iRow; k++)
    {
        ptKernelParams->ppppfKernelParams[k] = (float ***)malloc(sizeof(float **) * ptKernelParams->iCol);
        while (ptKernelParams->ppppfKernelParams[k] == NULL)
        {
            DBG_PRINTF(DLOG_LVL_ERROR,DLOG_TAG,"ppppfKernelParams[][%d] malloc  faild!\r\n", k);
            ptKernelParams->ppppfKernelParams[k] = (float ***)malloc(sizeof(float **) * ptKernelParams->iCol);
        }
        for (i = 0; i < ptKernelParams->iCol; i++)
        {
            ptKernelParams->ppppfKernelParams[k][i] = (float **)malloc(sizeof(float *) * ptKernelParams->iDim);
            while (ptKernelParams->ppppfKernelParams[k][i] == NULL)
            {
                DBG_PRINTF(DLOG_LVL_ERROR,DLOG_TAG,"ppppfKernelParams[][%d][%d] malloc  faild!\r\n", k, i);
                ptKernelParams->ppppfKernelParams[k][i] = (float **)malloc(sizeof(float *) * ptKernelParams->iDim);
            }
            for (j = 0; j < ptKernelParams->iDim; j++)
            {
                ptKernelParams->ppppfKernelParams[k][i][j] = (float *)malloc(sizeof(float) * ptKernelParams->iNum);
                while (ptKernelParams->ppppfKernelParams[k][i][j] == NULL)
                {
                    DBG_PRINTF(DLOG_LVL_ERROR,DLOG_TAG,"ppppfKernelParams[k][%d][%d][%d] malloc  faild!\r\n", k,i, j);
                    ptKernelParams->ppppfKernelParams[k][i][j] = (float *)malloc(sizeof(float) * ptKernelParams->iNum);
                }
            }
        }
    }
    ptKernelParams->pfKernelBias = (float *)malloc(sizeof(float) * ptKernelParams->iBiasLength);
}

/**
 * @brief 释放ptKernelParams内存
 * 
 * @param ptKernelParams 模型参数
 */
static void FreeForParams(PT_KernelParams ptKernelParams)
{
    int k,i,j;
    for (k = 0; k < ptKernelParams->iRow; k++)
    {
        for (i = 0; i < ptKernelParams->iCol; i++)
        {
            for (j = 0; j < ptKernelParams->iDim; j++)
            {
                free(ptKernelParams->ppppfKernelParams[k][i][j]);
            }
            free(ptKernelParams->ppppfKernelParams[k][i]);
        }
        free(ptKernelParams->ppppfKernelParams[k]);
    }
    free(ptKernelParams->ppppfKernelParams);
    ptKernelParams->ppppfKernelParams = NULL;
    free(ptKernelParams->pfKernelBias);
    ptKernelParams->pfKernelBias = NULL;
}

/**
 * @brief 打开模型文件
 * 
 * @param pcModelPath 模型参数文件路径
 * @return int 成功：0;失败：其他;
 */
static int H5Init(const char *pcModelPath) 
{
    // 打开HDF5文件
    g_File_id = H5Fopen(pcModelPath, H5F_ACC_RDONLY, H5P_DEFAULT);
    if (g_File_id < 0) {
        DBG_PRINTF(DLOG_LVL_ERROR,DLOG_TAG,"无法打开HDF5文件。\n");
        return -1;
    }
    return 0;
}

/**
 * @brief 将pfBuf中的内容存入ptKernelParams中
 * 
 * @param ptKernelParams 模型参数
 * @param pfBuf 从hd5文件中读取的参数
 */
static void GetKernelParamsFromBuf(PT_KernelParams ptKernelParams,float *pfBuf)
{
    uint32_t uiCurNum, uiCurDim, uiCurRow, uiCurCol;
    uint32_t uiTemp = 0;
    for (uiCurRow = 0; uiCurRow < ptKernelParams->iRow; uiCurRow++)
    {
        for (uiCurCol = 0; uiCurCol < ptKernelParams->iCol; uiCurCol++)
        {
            for (uiCurDim = 0; uiCurDim < ptKernelParams->iDim; uiCurDim++)
            {
                for (uiCurNum = 0; uiCurNum < ptKernelParams->iNum; uiCurNum++)
                {/*fprintf(txt_file, "%.8f ", data[i + j*dims[3] + k*(dims[1]*dims[2]*dims[3]) + l*(dims[2]*dims[3])]);*/
                    /*index = iCurRow *(C*D*N) + uiCurCol *(D*N) + uiCurDim *N + uiCurNum */
                    ptKernelParams->ppppfKernelParams[uiCurRow][uiCurCol][uiCurDim][uiCurNum] = pfBuf[uiTemp++];
                }
            }   
        }
    }
}

/**
 * @brief 获取模型参数
 * 
 * @param ptKernelParams 模型参数描述结构体
 * @return int 成功：0;失败：其他;
 */
static int GetKernelParams(PT_KernelParams ptKernelParams)
{
    hid_t DataSetID, DataSpaceID;
    herr_t Status;
    uint32_t iNum = 1;
    int iNumofDims;
    char strTemp[256];
    float *pfBuf;
    int i;
    sprintf(strTemp,"%s/%s",ptKernelParams->ParamsPathInH5,"kernel:0");
    strTemp[255] = '\0';
    // 打开数据集
    DataSetID = H5Dopen(g_File_id, strTemp, H5P_DEFAULT);
    if (DataSetID < 0) {
        DBG_PRINTF(DLOG_LVL_ERROR,DLOG_TAG,"无法打开数据集\n");
        H5Fclose(g_File_id);
        return -1;
    }
    // 获取数据集的数据空间
    DataSpaceID = H5Dget_space(DataSetID);
    if (DataSpaceID < 0) {
        DBG_PRINTF(DLOG_LVL_ERROR,DLOG_TAG,"无法获取数据空间\n");
        H5Dclose(DataSpaceID);
        return -1;
    }
    // 获取数据集的维度数量
    iNumofDims = H5Sget_simple_extent_ndims(DataSpaceID);
    if (iNumofDims < 0) {
        DBG_PRINTF(DLOG_LVL_ERROR,DLOG_TAG,"无法获取维度数量\n");
        H5Sclose(DataSpaceID);
        H5Dclose(DataSetID);
        return -1;
    }
    // 获取数据集的维度信息
    hsize_t dims[iNumofDims];
    Status = H5Sget_simple_extent_dims(DataSpaceID, dims, NULL);
    if (Status < 0) {
        DBG_PRINTF(DLOG_LVL_ERROR,DLOG_TAG,"无法获取维度信息\n");
        H5Sclose(DataSpaceID);
        H5Dclose(DataSetID);
        return -1;
    }
    for(i =0;i<iNumofDims;i++)
    {
        iNum *= dims[i];
    }
    //DBG_PRINTF(DLOG_LVL_ERROR,DLOG_TAG,"iNum: %lu %s %d\r\n",iNum,__FILE__,__LINE__);
    pfBuf = (float *)malloc(sizeof(float)*iNum);
    //读取数据集内容到内存中的数组
    Status = H5Dread(DataSetID, H5T_NATIVE_FLOAT, H5S_ALL, H5S_ALL, H5P_DEFAULT,pfBuf);
    if (Status < 0) {
        DBG_PRINTF(DLOG_LVL_ERROR,DLOG_TAG,"无法读取数据集\n");
        H5Dclose(DataSetID);
        return -1;
    }
      
    GetKernelParamsFromBuf(ptKernelParams,pfBuf);
     
    free(pfBuf);
    H5Dclose(DataSetID);
    return 0;
}

/**
 * @brief 读取模型偏移量
 * 
 * @param ptKernelParams 模型参数描述结构体
 * @return int 成功：0;失败：其他;
 */
static int GetBiasParams(PT_KernelParams ptKernelParams)
{
    hid_t DataSetID;
    herr_t Status;
    float *pfBuf;
    char strTemp[256];

    sprintf(strTemp,"%s/%s",ptKernelParams->ParamsPathInH5,"bias:0");
    strTemp[255] = '\0';
    // 打开数据集
    DataSetID = H5Dopen(g_File_id,strTemp, H5P_DEFAULT);
    if (DataSetID < 0) {
        DBG_PRINTF(DLOG_LVL_ERROR,DLOG_TAG,"无法打开数据集\n");
        H5Fclose(g_File_id);
        return -1;
    }
    // 读取数据集内容到内存中的数组
    pfBuf = (float *)malloc(sizeof(float)*ptKernelParams->iBiasLength);
    Status = H5Dread(DataSetID, H5T_NATIVE_FLOAT, H5S_ALL, H5S_ALL, H5P_DEFAULT, pfBuf);
    if (Status < 0) {
        DBG_PRINTF(DLOG_LVL_ERROR,DLOG_TAG,"无法读取数据集\n");
        H5Dclose(DataSetID);
        return -1;
    }
    memcpy(ptKernelParams->pfKernelBias,pfBuf,sizeof(float)*ptKernelParams->iBiasLength);
    free(pfBuf);
    H5Dclose(DataSetID);
    return 0;
}

// PT_KernelParams GetParamsDisc()
// {
//     return g_tKernelParams;
// }

/**
 * @brief 初始化模型参数
 * 
 * @param ptKernelParams 模型参数描述结构体
 * @param iNumOfLayer 模型层数
 * @param pcModelPath 模型参数路径
 * @return int 
 */
int ModelParamsInit(PT_KernelParams ptKernelParams,int iNumOfLayer,char* pcModelPath)
{
    int iError,i;
    iError = H5Init(pcModelPath);
    if(iError < 0)
    {
        return iError;
    } 
    for (i = 0; i < iNumOfLayer; i++)
    {
        //DBG_PRINTF(DLOG_LVL_ERROR,DLOG_TAG,"I : %d\r\n",i);
        MallocForParams(&ptKernelParams[i]);
        iError = GetKernelParams(&ptKernelParams[i]);
        if (iError < 0)
        {
            return iError;
        }
          
        iError = GetBiasParams(&ptKernelParams[i]);
        if (iError < 0)
        {
            return iError;
        }    
    }
    H5Fclose(g_File_id);
    return iError;
}

/**
 * @brief 释放模型参数
 * 
 * @param ptKernelParams 模型参数
 * @param iNumOfLayer 模型层数
 */
void ModelParamsFree(PT_KernelParams ptKernelParams,int iNumOfLayer)
{
    int i;
    for (i = 0; i < iNumOfLayer; i++)
    {
        FreeForParams(&ptKernelParams[i]);  
    }
}